import os
import sys
import math
from time import time

import torch
import torch.nn.functional as F

def accuracy(logit, target, topk=(1,)):
    """Computes the precision@k for the specified values of k"""
    output = F.softmax(logit, dim=1)
    maxk = max(topk)
    batch_size = target.size(0)

    _, pred = output.topk(maxk, 1, True, True)
    pred = pred.t()
    correct = pred.eq(target.view(1, -1).expand_as(pred))

    res = []
    for k in topk:
        correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
        res.append(correct_k.mul_(100.0 / batch_size))
    return res

# Evaluate the Model
def evaluate(test_loader, model, device):
    model.eval()    # Change model to 'eval' mode.
    correct = 0
    loss = 0
    total = 0
    with torch.no_grad():
        for images, labels, _, _ in test_loader:
            images = images.to(device)
            labels = labels.to(device)
            logits = model(images)
            outputs = F.softmax(logits, dim=1)
            _, pred = torch.max(outputs.data, 1)
            loss += F.cross_entropy(logits, labels)
            total += labels.size(0)
            correct += (pred.cpu() == labels.cpu()).sum()
            #print('correct:', correct)
            #print('total:', total)

    acc = 100 * correct / total
    loss /= total
    return acc, loss
